//
// Created by ubuntu on 8/19/25.
//

#include "H264Utils.h"

#include <cmath>

unsigned int H264Utils::ue(const unsigned char *pBuff, const unsigned int nLen, unsigned int &nStartBit) {
    // 计算前导零比特的数量
    unsigned int nZeroNum = 0;
    while (nStartBit < nLen * 8) {
        // 检查当前位是否为1，0x80是二进制10000000
        if (pBuff[nStartBit / 8] & (0x80 >> (nStartBit % 8))) {
            break; // 找到第一个1则停止计数
        }
        nZeroNum++; // 增加零比特计数
        nStartBit++; // 移动到下一个比特
    }
    nStartBit++; // 跳过值为1的比特位

    // 根据前导零比特数量解码后续比特
    unsigned long dwRet = 0;
    for (unsigned int i = 0; i < nZeroNum; i++) {
        dwRet <<= 1; // 左移一位，为新比特腾出位置
        // 检查当前位是否为1
        if (pBuff[nStartBit / 8] & (0x80 >> (nStartBit % 8))) {
            dwRet += 1; // 如果是1则加1
        }
        nStartBit++; // 移动到下一个比特
    }
    // 指数哥伦布编码解码公式: 2^(leadingZeroBits) - 1 + readBits(leadingZeroBits)
    return (1 << nZeroNum) - 1 + dwRet;
}

/**
* 解码SE(Exp-Golomb)
*
* @param pBuff SE数据内容
* @param nLen SE数据的长度
* @param nStartBit 当前解析到的字节位置
*
* @返回码 解码后的有符号整数值
*/
int H264Utils::se(const unsigned char *pBuff, const unsigned int nLen, unsigned int &nStartBit) {
    // 先按无符号指数哥伦布编码方式解码
    const int UeVal = ue(pBuff, nLen, nStartBit);
    double k = UeVal;
    // 使用公式 (k+1)/2 向上取整得到绝对值
    int nValue = std::ceil(k / 2);
    // 如果ue值为偶数，则结果为负数
    if (UeVal % 2 == 0) {
        nValue = -nValue;
    }
    return nValue;
}

/**
 * 解码u(BitCount)
 * @param BitCount 位数
 * @param buf 输入数据
 * @param nStartBit 当前解析到的字节位置
 * @return 解码后的无符号整数值
 */
unsigned long H264Utils::u(const unsigned int BitCount, const unsigned char *buf, unsigned int &nStartBit) {
    unsigned long dwRet = 0;
    // 逐位读取指定数量的比特
    for (unsigned int i = 0; i < BitCount; i++) {
        dwRet <<= 1; // 左移一位
        // 检查当前位是否为1
        if (buf[nStartBit / 8] & (0x80 >> (nStartBit % 8))) {
            dwRet += 1; // 如果是1则加1
        }
        nStartBit++; // 移动到下一个比特
    }
    return dwRet;
}

/**
 * H264的NAL起始码防竞争机制
 *
 * @param buf SPS数据内容
 * @param buf_size SPS数据长度
 *
 * @无返回值
 */
void H264Utils::deEmulationPrevention(unsigned char *buf, unsigned int *buf_size) {
    int i = 0, j = 0;
    unsigned char *tmp_ptr = nullptr;
    unsigned int tmp_buf_size = 0;
    int val = 0;

    tmp_ptr = buf;
    tmp_buf_size = *buf_size;
    // 遍历缓冲区查找并移除防止竞争的字节0x03
    for (i = 0; i < (tmp_buf_size - 2); i++) {
        // 检查是否存在0x000003模式
        val = (tmp_ptr[i] ^ 0x00) + (tmp_ptr[i + 1] ^ 0x00) + (tmp_ptr[i + 2] ^ 0x03);
        if (val == 0) {
            // 移除0x03字节，将后续字节前移
            for (j = i + 2; j < tmp_buf_size - 1; j++)
                tmp_ptr[j] = tmp_ptr[j + 1];

            // 减少缓冲区大小
            (*buf_size)--;
        }
    }

    return;
}

/**
 * 解码SPS,获取视频图像宽、高信息
 *
 * @param buf SPS数据内容
 * @param nLen SPS数据的长度
 * @param width 图像宽度
 * @param height 图像高度
 * @param fps 帧率

 * @成功则返回1 , 失败则返回0
 */
int H264Utils::h264DecodeSps(unsigned char *buf, unsigned int nLen, int &width, int &height, int &fps) {
    unsigned int startBit = 0; // 初始化比特读取位置
    fps = 0; // 初始化帧率为0
    deEmulationPrevention(buf, &nLen); // 处理防止竞争的字节

    // 解析NAL头信息
    u(1, buf, startBit); // forbidden_zero_bit (1位)
    u(2, buf, startBit); // nal_ref_idc (2位)
    const int nalUnitType = u(5, buf, startBit); //nal_unit_type (5位)

    //非SPS数据
    if (nalUnitType != 7) {
        return false; // 如果不是SPS NAL单元，返回失败
    }

    //解码SPS 数据
    const int profileIdc = u(8, buf, startBit); //profile_idc (8位)
    u(1, buf, startBit); // constraint_set0_flag (1位)
    u(1, buf, startBit); // constraint_set1_flag (1位)
    u(1, buf, startBit); // constraint_set2_flag (1位)
    u(1, buf, startBit); // constraint_set3_flag (1位)
    u(4, buf, startBit); // reserved_zero_4bits (4位)
    u(8, buf, startBit); // level_idc (8位)

    ue(buf, nLen, startBit); // seq_parameter_set_id (ue编码)

    //跳过seq_parameter_set_id
    if (profileIdc == 100 || profileIdc == 110 ||
        profileIdc == 122 || profileIdc == 144) {
        const int chroma_format_idc = ue(buf, nLen, startBit); // 色度格式
        if (chroma_format_idc == 3) {
            u(1, buf, startBit); // separate_colour_plane_flag
        }
        ue(buf, nLen, startBit); // bit_depth_luma_minus8
        ue(buf, nLen, startBit); // bit_depth_chroma_minus8
        u(1, buf, startBit); // qpprime_y_zero_transform_bypass_flag
        const int seqScalingMatrixPresentFlag = u(1, buf, startBit); //seq_scaling_matrix_present_flag
        if (seqScalingMatrixPresentFlag) {
            // 处理缩放矩阵标志
            for (int i = 0; i < 8; i++) {
                u(1, buf, startBit); // seq_scaling_list_present_flag[i]
            }
        }
    }
    ue(buf, nLen, startBit); // log2_max_frame_num_minus4
    const int pic_order_cnt_type = ue(buf, nLen, startBit); // pic_order_cnt_type
    if (pic_order_cnt_type == 0) {
        ue(buf, nLen, startBit); // log2_max_pic_order_cnt_lsb_minus4
    } else if (pic_order_cnt_type == 1) {
        u(1, buf, startBit); // delta_pic_order_always_zero_flag
        se(buf, nLen, startBit); // offset_for_non_ref_pic
        se(buf, nLen, startBit); // offset_for_top_to_bottom_field
        const int num_ref_frames_in_pic_order_cnt_cycle = ue(buf, nLen, startBit);
        // num_ref_frames_in_pic_order_cnt_cycle

        // 分配并解析参考帧偏移数组
        const auto offset_for_ref_frame = new int[num_ref_frames_in_pic_order_cnt_cycle];
        for (int i = 0; i < num_ref_frames_in_pic_order_cnt_cycle; i++) {
            offset_for_ref_frame[i] = se(buf, nLen, startBit); // offset_for_ref_frame[i]
        }
        delete [] offset_for_ref_frame; // 释放内存
    }
    ue(buf, nLen, startBit); // max_num_ref_frames
    u(1, buf, startBit); // gaps_in_frame_num_value_allowed_flag
    const int pic_width_in_mbs_minus1 = ue(buf, nLen, startBit); // 图像宽度(以宏块为单位)
    const int pic_height_in_map_units_minus1 = ue(buf, nLen, startBit); // 图像高度(以宏块为单位)

    // 计算实际图像尺寸(每个宏块为16x16像素)
    width = (pic_width_in_mbs_minus1 + 1) * 16;
    height = (pic_height_in_map_units_minus1 + 1) * 16;

    const int frame_mbs_only_flag = u(1, buf, startBit); // frame_mbs_only_flag
    if (!frame_mbs_only_flag) {
        u(1, buf, startBit); // mb_adaptive_frame_field_flag
    }

    u(1, buf, startBit); // direct_8x8_inference_flag
    const int frame_cropping_flag = u(1, buf, startBit); // frame_cropping_flag
    if (frame_cropping_flag) {
        // 处理帧裁剪偏移
        ue(buf, nLen, startBit); // frame_crop_left_offset
        ue(buf, nLen, startBit); // frame_crop_right_offset
        ue(buf, nLen, startBit); // frame_crop_top_offset
        ue(buf, nLen, startBit); // frame_crop_bottom_offset
    }
    const int vui_parameter_present_flag = u(1, buf, startBit); // vui_parameter_present_flag
    if (vui_parameter_present_flag) {
        const int aspect_ratio_info_present_flag = u(1, buf, startBit); // aspect_ratio_info_present_flag
        if (aspect_ratio_info_present_flag) {
            const int aspect_ratio_idc = u(8, buf, startBit); // aspect_ratio_idc
            // 处理扩展纵横比信息
            if (aspect_ratio_idc == 255) {
                u(16, buf, startBit); // sar_width
                u(16, buf, startBit); // sar_height
            }
        }
        const int overscan_info_present_flag = u(1, buf, startBit); // overscan_info_present_flag
        if (overscan_info_present_flag) {
            u(1, buf, startBit); // overscan_appropriate_flag
        }
        const int video_signal_type_present_flag = u(1, buf, startBit); // video_signal_type_present_flag
        if (video_signal_type_present_flag) {
            u(3, buf, startBit); // video_format
            u(1, buf, startBit); // video_full_range_flag
            const int colour_description_present_flag = u(1, buf, startBit); // colour_description_present_flag
            if (colour_description_present_flag) {
                u(8, buf, startBit); // colour_primaries
                u(8, buf, startBit); // transfer_characteristics
                u(8, buf, startBit); // matrix_coefficients
            }
        }
        const int chroma_loc_info_present_flag = u(1, buf, startBit); // chroma_loc_info_present_flag
        if (chroma_loc_info_present_flag) {
            ue(buf, nLen, startBit); // chroma_sample_loc_type_top_field
            ue(buf, nLen, startBit); // chroma_sample_loc_type_bottom_field
        }
        const int timing_info_present_flag = u(1, buf, startBit); // timing_info_present_flag
        if (timing_info_present_flag) {
            const int num_units_in_tick = u(32, buf, startBit); // num_units_in_tick
            const int time_scale = u(32, buf, startBit); // time_scale
            // 计算帧率: time_scale / (2 * num_units_in_tick)
            fps = time_scale / (2 * num_units_in_tick);
        }
    }

    return true; // 解码成功
}
